#include "symbols.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/ptrace.h>
#include <sys/wait.h>
#include <fcntl.h>
#include <errno.h>
#include <limits.h>
#include <sys/stat.h>
#include <libgen.h>
#include <elf.h>

// 如果PATH_MAX未定义，定义一个默认值
#ifndef PATH_MAX
#define PATH_MAX 4096
#endif

#define SYMBOL_TABLE_INIT_CAPACITY 64

// 创建符号表
symbol_table_t* symbol_table_create(void) {
    symbol_table_t *table = malloc(sizeof(symbol_table_t));
    if (!table) return NULL;
    
    table->count = 0;
    table->capacity = SYMBOL_TABLE_INIT_CAPACITY;
    table->symbols = malloc(sizeof(symbol_info_t) * table->capacity);
    
    if (!table->symbols) {
        free(table);
        return NULL;
    }
    
    return table;
}

// 销毁符号表
void symbol_table_destroy(symbol_table_t *table) {
    if (table) {
        if (table->symbols) {
            free(table->symbols);
        }
        free(table);
    }
}

unsigned long get_text_segment_offset(const char *filename) {
    int fd = open(filename, O_RDONLY);
    if (fd < 0) return 0;
    
    Elf64_Ehdr ehdr;
    if (read(fd, &ehdr, sizeof(ehdr)) != sizeof(ehdr)) {
        close(fd);
        return 0;
    }
    
    // 读取程序头表
    Elf64_Phdr *phdr = (Elf64_Phdr *)malloc(ehdr.e_phentsize * ehdr.e_phnum);
    lseek(fd, ehdr.e_phoff, SEEK_SET);
    if (read(fd, phdr, ehdr.e_phentsize * ehdr.e_phnum) != ehdr.e_phentsize * ehdr.e_phnum) {
        free(phdr);
        close(fd);
        return 0;
    }
    
    unsigned long text_offset = 0;
    for (int i = 0; i < ehdr.e_phnum; i++) {
        if (phdr[i].p_type == PT_LOAD && (phdr[i].p_flags & PF_X)) {
            text_offset = phdr[i].p_offset;
            break;
        }
    }
    
    free(phdr);
    close(fd);
    return text_offset;
}

// 执行命令并获取输出
static char* execute_command(const char *command) {
    FILE *pipe = popen(command, "r");
    if (!pipe) return NULL;
    
    char buffer[128];
    char *result = (char*)malloc(1024*sizeof(char));
    if (!result) {
        pclose(pipe);
        return NULL;
    }
    
    result[0] = '\0';
    
    while (fgets(buffer, sizeof(buffer), pipe) != NULL) {
        strcat(result, buffer);
    }

    pclose(pipe);
    return result;
}



// 从 /proc/pid/maps 中精确获取目标程序的代码段基地址
unsigned long get_exact_code_base(pid_t pid, const char *target_exe) {
    char maps_path[256];
    FILE *maps_file;
    char line[1024];
    unsigned long code_base = 0;
    char real_path[1024];
    
    // 获取目标程序的真实路径
    if (realpath(target_exe, real_path) == NULL) {
        strncpy(real_path, target_exe, sizeof(real_path) - 1);
    }
    
    snprintf(maps_path, sizeof(maps_path), "/proc/%d/maps", pid);
    
    maps_file = fopen(maps_path, "r");
    if (!maps_file) {
        perror("fopen maps");
        return 0;
    }
    
    printf("Searching for: %s\n", real_path);
    
    while (fgets(line, sizeof(line), maps_file)) {
        unsigned long start, end;
        char perms[5];
        char path[256] = "";
        unsigned long offset;
        
        // 解析 maps 行：start-end perms offset dev inode path
        if (sscanf(line, "%lx-%lx %4s %lx", &start, &end, perms, &offset) >= 4) {
            // 检查权限是否为可执行
            if (strstr(perms, "x") && strstr(perms, "r")) {
                // 提取路径
                char *path_start = strchr(line, '/');
                if (path_start) {
                    char *newline = strchr(path_start, '\n');
                    if (newline) *newline = '\0';
                    strncpy(path, path_start, sizeof(path) - 1);
                    
                    // 精确匹配路径
                    if (strcmp(path, real_path) == 0) {
                        printf("Found exact match: %s at 0x%lx (offset: 0x%lx)\n", path, start, offset);
                        code_base = start;
                        break;
                    }
                }
            }
        }
    }
    
    fclose(maps_file);
    return code_base;
}

// 使用nm获取符号偏移（支持函数和全局变量）
static uintptr_t get_function_offset_using_nm(const char *filename, const char *symbol_name) {
    char command[512];
    char *output;
    uintptr_t offset = 0;
    char line[256];
    FILE *fp;
    
    // 支持全局变量符号（D表示数据段，B表示BSS段，T表示代码段）
    snprintf(command, sizeof(command), "nm '%s' 2>/dev/null | grep ' %s$'", filename, symbol_name);
    printf("执行命令: %s\n", command);
    fp = popen(command, "r");
    if (fp == NULL) {
        perror("popen failed");
        return 0;
    }
    
    while (fgets(line, sizeof(line), fp) != NULL) {
        printf("命令输出:\n%s\n", line);
        // nm 输出格式: 0000000000001149 T main (函数)
        // nm 输出格式: 0000000000004050 D global_array_var (全局变量)
        char symbol_type;
        if (sscanf(line, "%lx %c", &offset, &symbol_type) == 2) {
            // 检查符号类型：T/t (函数), D/d (数据段), B/b (BSS段)
            if (symbol_type == 'T' || symbol_type == 't' || 
                symbol_type == 'D' || symbol_type == 'd' ||
                symbol_type == 'B' || symbol_type == 'b') {
                printf("Found using nm: %s (type: %c) at offset 0x%lx\n", symbol_name, symbol_type, offset);
                pclose(fp);
                return offset;
            }
        }
    }
    pclose(fp);
    
    return 0;
}

// 使用objdump获取函数偏移
static uintptr_t get_function_offset_using_objdump(const char *filename, const char *func_name) {
    char command[512];
    char *output;
    uintptr_t offset = 0;
    
    snprintf(command, sizeof(command), "objdump -t '%s' 2>/dev/null | grep ' %s$'", filename, func_name);
    output = execute_command(command);
    
    if (output) {
        char *line = strtok(output, "\n");
        while (line) {
            // objdump 输出格式: 0000000000001149 g     F .text  0000000000000018 main
            uintptr_t value;
            char name[256];
            
            if (sscanf(line, "%lx %*s %*s %*s %*s %255s", &value, name) == 2) {
                if (strcmp(name, func_name) == 0) {
                    offset = value;
                    break;
                }
            }
            line = strtok(NULL, "\n");
        }
        free(output);
    }
    
    return offset;
}

// 从程序文件加载符号表
int symbol_table_load_from_program(symbol_table_t *table, const char *program_path) {
    if (!table || !program_path) return -1;
    
    // 使用nm获取符号信息
    char command[512];
    snprintf(command, sizeof(command), "nm -n '%s' 2>/dev/null", program_path);
    char *output = execute_command(command);
    
    if (!output) {
        // 尝试使用objdump
        snprintf(command, sizeof(command), "objdump -t '%s' 2>/dev/null", program_path);
        output = execute_command(command);
        
        if (!output) {
            return -1; // 无法获取符号信息
        }
    }
    
    char *line = strtok(output, "\n");
    while (line) {
        // 解析符号行
        uintptr_t address;
        char type;
        char name[256];
        
        if (sscanf(line, "%lx %c %255s", &address, &type, name) == 3) {
            // 只处理函数符号（T/t）
            if (type == 'T' || type == 't') {
                // 扩展符号表容量（如果需要）
                if (table->count >= table->capacity) {
                    size_t new_capacity = table->capacity * 2;
                    symbol_info_t *new_symbols = realloc(table->symbols, 
                                                         sizeof(symbol_info_t) * new_capacity);
                    if (!new_symbols) {
                        free(output);
                        return -1;
                    }
                    table->symbols = new_symbols;
                    table->capacity = new_capacity;
                }
                
                // 添加符号信息
                symbol_info_t *sym = &table->symbols[table->count];
                strncpy(sym->name, name, sizeof(sym->name) - 1);
                sym->address = address;
                sym->filename[0] = '\0';
                sym->line = 0;
                table->count++;
            }
        }
        line = strtok(NULL, "\n");
    }
    
    free(output);
    return 0;
}

// 在符号表中查找地址
uintptr_t symbol_table_find_address(symbol_table_t *table, const char *symbol_name) {
    if (!table || !symbol_name) return 0;
    
    for (size_t i = 0; i < table->count; i++) {
        if (strcmp(table->symbols[i].name, symbol_name) == 0) {
            return table->symbols[i].address;
        }
    }
    
    return 0;
}

// 根据地址查找符号名
int symbol_table_find_symbol_by_address(symbol_table_t *table, uintptr_t address, 
                                       char *symbol_name, size_t name_size) {
    if (!table || !symbol_name) return -1;
    
    for (size_t i = 0; i < table->count; i++) {
        if (table->symbols[i].address == address) {
            strncpy(symbol_name, table->symbols[i].name, name_size - 1);
            symbol_name[name_size - 1] = '\0';
            return 0;
        }
    }
    
    return -1;
}

// 在进程中查找符号地址（支持函数和全局变量）
uintptr_t find_function_address_in_process(pid_t pid, const char *program_path, const char *symbol_name) {
    if (pid == 0 || !program_path || !symbol_name) return 0;
    
    // 步骤1: 获取符号在文件中的偏移
    uintptr_t symbol_offset = get_function_offset_using_nm(program_path, symbol_name);
    if (!symbol_offset) {
        symbol_offset = get_function_offset_using_objdump(program_path, symbol_name);
        if (!symbol_offset) {
            printf("无法找到符号 '%s' 在文件中的偏移\n", symbol_name);
            return 0;
        }
    }
    
    printf("符号 '%s' 在文件中的偏移: 0x%lx\n", symbol_name, symbol_offset);
    
    // 步骤2: 获取进程中的代码段基地址
    uintptr_t code_base = get_exact_code_base(pid, program_path);
    if (!code_base) {
        printf("无法获取进程 %d 的代码段基地址\n", pid);
        return 0;
    }
    
    printf("进程 %d 的代码段基地址: 0x%lx\n", pid, code_base);
    
    // 步骤3: 计算运行时地址
    // 根据DWARF信息，变量偏移是相对于程序基地址的
    // 使用/proc/pid/maps的第一行作为基地址
    uintptr_t runtime_addr = 0;
    
    // 从/proc/pid/maps中读取第一行作为基地址
    char maps_path[256];
    snprintf(maps_path, sizeof(maps_path), "/proc/%d/maps", pid);
    FILE *maps_file = fopen(maps_path, "r");
    if (maps_file) {
        char line[1024];
        if (fgets(line, sizeof(line), maps_file)) {
            printf("line: %s", line);
            unsigned long start, end, offset;
            char perms[5];
            char path[256] = "";
            
            if (sscanf(line, "%lx-%lx %4s %lx", &start, &end, perms, &offset) >= 4) {
                // 使用第一行的起始地址作为基地址
                // 根据DWARF信息：基地址 + 变量偏移 = 运行时地址
                runtime_addr = start + symbol_offset;
                printf("使用第一行映射作为基地址: 0x%lx\n", start);
                printf("符号 '%s' 在进程中的运行时地址: 0x%lx\n", symbol_name, runtime_addr);
            }
        }
        fclose(maps_file);
    }
    
    // 如果第一行查找失败，回退到代码段计算
    if (runtime_addr == 0) {
        unsigned long text_offset = get_text_segment_offset(program_path);
        runtime_addr = code_base + (symbol_offset - text_offset);
        printf("符号 '%s' 在进程中的运行时地址（回退计算）: 0x%lx\n", symbol_name, runtime_addr);
    }
    
    // 步骤4: 验证地址是否有效
    errno = 0;
    long data = ptrace(PTRACE_PEEKDATA, pid, (void*)runtime_addr, NULL);
    if (data == -1 && errno != 0) {
        printf("警告: 地址 0x%lx 可能无效 (ptrace失败)\n", runtime_addr);
        // 但仍然返回地址，让调用者决定是否使用
    } else {
        printf("地址验证成功，读取到的数据: 0x%lx\n", data);
    }
    
    return runtime_addr;
}